# app/models/schemas.py
from pydantic import BaseModel
from typing import List, Optional, Dict, Literal


class ModelRequest(BaseModel):
    prompt: str
    thinking: Optional[bool] = False


class SimpleQuery(BaseModel):
    keyword: str
    topK: Optional[int] = 3
    collectionName: Optional[str] = "my_knowledge"


class Query(BaseModel):
    keyword: str
    topK: int
    collectionName: str
    searchMethod: Literal["dense", "sparse", "hybrid"] = "hybrid"
    rerankWeight: Optional[float] = 0.7
    rerankMethod: Literal["default", "bge", "rrf"] = "bge"
    scoreThreshold: Optional[float] = 0.5
    thinking: Optional[bool] = False


class Rerank(BaseModel):
    query: str
    documents: List[str]


class Document(BaseModel):
    text: str
    metadata: Optional[Dict] = None
    collection_name: Optional[str] = "my_knowledge"


class DocumentBatch(BaseModel):
    documents: List[Document]
